import warnings
warnings.filterwarnings('ignore')
import util_for_output_zh

import pandas as pd
import numpy as np
# import matplotlib.pyplot as plt
from util_set_zh_matplot import plt
import seaborn as sns
from sklearn.ensemble import RandomForestRegressor
from sklearn.preprocessing import PolynomialFeatures, StandardScaler
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_absolute_error, mean_squared_error
from skopt import gp_minimize
from skopt.space import Real
from typing import Tuple, List, Dict
from tqdm import tqdm
from joblib import Parallel, delayed  # 并行计算
import multiprocessing  # 多进程
from pyswarm import pso

import os


# 设置显示选项
pd.set_option('display.max_columns', None)
pd.set_option('display.max_rows', 100)
pd.set_option('display.width', None)
pd.set_option('display.max_colwidth', None)

import numpy as np
from scipy.stats import norm

out_dir = 'ques3_wucha'
os.makedirs(out_dir,exist_ok=True)

'''
额外的约束条件 需要 calculate_pass_probability(...) > 0.8
'''
def calculate_y_chromosome(身高, 体重, 年龄, BMI, 孕周):
    """
    根据三次项多项式公式计算Y染色体浓度
    
    参数:
        身高: 数值或数组，单位cm
        体重: 数值或数组，单位kg
        年龄: 数值或数组，单位岁
        BMI: 数值或数组，身体质量指数
        孕周: 数值或数组，检测时的孕周
        
    返回:
        数值或数组，计算得到的Y染色体浓度
    """
    # 计算各变量的一次项、二次项和三次项
    浓度 = (
        # 常数项
        0.090230 +
        
        # 一次项
        -0.450695 * 身高 +
        0.726708 * 体重 +
        -0.002885 * 年龄 +
        -0.548887 * BMI +
        -0.009580 * 孕周 +
        
        # 二次项
        0.005772 * (身高 **2) +
        -0.039292 * (体重** 2) +
        0.000414 * (年龄 **2) +
        0.019534 * (BMI** 2) +
        -0.005537 * (孕周 **2) +
        
        # 三次项
        0.000265 * (身高** 3) +
        0.001969 * (体重 **3) +
        -0.000182 * (年龄** 3) +
        -0.000615 * (BMI **3) +
        0.007666 * (孕周** 3)
    )
    return 浓度

def calculate_pass_probability(身高, 体重, 年龄, BMI, 孕周, 预测误差标准差):
    """
    计算Y染色体浓度达标的概率 (P(Yi ≥ 0.04))
    
    参数:
        身高: 数值或数组，单位cm
        体重: 数值或数组，单位kg
        年龄: 数值或数组，单位岁
        BMI: 数值或数组，身体质量指数
        孕周: 数值或数组，检测时的孕周
        预测误差标准差: 模型预测的标准差
        
    返回:
        数值或数组，达标的概率
    """
    # 计算预测浓度
    预测浓度 = calculate_y_chromosome(身高, 体重, 年龄, BMI, 孕周)
    
    # 计算达标概率: P(Yi ≥ 0.04) = 1 - Φ((0.04 - ÛYi)/σi)
    z_score = (0.04 - 预测浓度) / 预测误差标准差
    达标概率 = 1 - norm.cdf(z_score)
    
    return 达标概率


# ---------------------------
# 增强版数据预处理
# ---------------------------
def preprocess_data(df: pd.DataFrame) -> pd.DataFrame:
    """处理原始数据，提取男胎样本并计算达标时间
    
    Args:
        df: 包含NIPT检测数据的DataFrame
        
    Returns:
        处理后的DataFrame，仅包含男胎样本和达标时间
    """
    # 检查必要列是否存在
    required_columns = ['孕妇代码', 'Y染色体浓度', '检测孕周']
    if not all(col in df.columns for col in required_columns):
        raise ValueError("输入数据缺少必要的列：'孕妇代码', 'Y染色体浓度', '检测孕周'")
    
    # 增强版孕周转换函数
    def convert_gestational_week(week_str):
        try:
            if isinstance(week_str, (int, float)):
                return float(week_str)
            if '+' in week_str:
                w, d = week_str.lower().replace('w', '').split('+')
                return float(w) + float(d)/7
            return float(week_str.lower().replace('w', ''))
        except:
            return np.nan
    
    df['检测孕周'] = df['检测孕周'].apply(convert_gestational_week)
    df = df.dropna(subset=['检测孕周'])
    
    # 筛选男胎数据（Y染色体浓度≥4%为男胎）
    male_data = df[df['Y染色体浓度'] >= 0.04].copy()
    
    # 更高效地计算达标孕周
    male_data['达标孕周'] = male_data.groupby('孕妇代码')['检测孕周'].transform('min')
    male_data = male_data.dropna(subset=['达标孕周'])
    
    # 只保留10-25周的数据，符合检测时间范围要求
    male_data = male_data[(male_data['检测孕周'] >= 10) & (male_data['检测孕周'] <= 25)]
    
    return male_data

# ---------------------------
# 增强版预测模型构建
# ---------------------------
def build_time_prediction_model(data: pd.DataFrame) -> Tuple[RandomForestRegressor, PolynomialFeatures, StandardScaler]:
    """建立达标时间预测模型
    
    Args:
        data: 预处理后的数据
        
    Returns:
        model: 训练好的随机森林模型
        poly: 多项式特征转换器
        scaler: 标准化转换器
    """
    # 检查必要特征列是否存在
    features = ['年龄', '身高', '体重', '孕妇BMI', '怀孕次数', '生产次数']
    missing_features = [f for f in features if f not in data.columns]
    if missing_features:
        raise ValueError(f"输入数据缺少必要的特征列：{missing_features}")
    
    X = data[features]
    y = data['达标孕周']
    
    # 添加多项式特征和标准化
    poly = PolynomialFeatures(degree=2, include_bias=False)
    scaler = StandardScaler()
    X_poly = scaler.fit_transform(poly.fit_transform(X))
    
    # 训练随机森林模型
    X_train, X_test, y_train, y_test = train_test_split(
        X_poly, y, test_size=0.2, random_state=42
    )
    
    model = RandomForestRegressor(
        n_estimators=150,
        max_depth=12,
        min_samples_split=5,
        random_state=42,
        n_jobs=-1
    )
    model.fit(X_train, y_train)
    
    # 更全面的模型评估
    train_pred = model.predict(X_train)
    test_pred = model.predict(X_test)
    
    print("\n=== 模型评估 ===")
    print(f"训练集R²: {model.score(X_train, y_train):.3f}")
    print(f"测试集R²: {model.score(X_test, y_test):.3f}")
    print(f"训练集MAE: {mean_absolute_error(y_train, train_pred):.3f} 周")
    print(f"测试集MAE: {mean_absolute_error(y_test, test_pred):.3f} 周")
    print(f"训练集RMSE: {np.sqrt(mean_squared_error(y_train, train_pred)):.3f} 周")
    print(f"测试集RMSE: {np.sqrt(mean_squared_error(y_test, test_pred)):.3f} 周")
    
    return model, poly, scaler

# ---------------------------
# 改进的评分计算函数（优先考虑更早孕周）
# ---------------------------
def get_time_weight(week: float) -> float:
    """增强对更早孕周的权重，鼓励提前检测"""
    if week < 11:
        return 2.5  # 提高权重，更鼓励早检测
    elif 11 <= week <= 12:
        return 2.0   # 提高权重
    elif 12 < week <= 15:
        return 1.7   # 提高权重
    elif 15 < week <= 17:
        return 1.5
    elif 17 < week <= 19:
        return 1.3
    elif 19 < week <= 21:
        return 1.1
    elif 21 < week <= 23:
        return 0.9   # 降低权重，不鼓励过晚检测
    elif 23 < week <= 25:
        return 0.7   # 降低权重
    else:
        return 0.0

def calculate_group_score(group_data: pd.DataFrame, t: float) -> float:
    """改进的评分函数，增加对更早孕周的偏好"""
    t_value = t[0] if isinstance(t, (list, np.ndarray)) else t
    
    # 对过晚的检测施加更强惩罚
    if t_value > 23:
        return 0.0  # 不考虑23周以后的检测
    
    valid_samples = group_data[group_data['Y染色体浓度'] > 0.04]
    if len(valid_samples) == 0:
        return 0.0
        
    # 计算达标比例
    early_records = valid_samples[valid_samples['检测孕周'] <= t_value]
    达标比例 = len(early_records) / len(valid_samples)
    
    # 获取时间权重
    time_weight = get_time_weight(t_value)
    
    # 关键改进：对接近的分数，更早的孕周有优势
    # 加入时间因子，让更早的孕周在分数相近时有更高优先级
    time_factor = 1.0 + (20 - t_value) * 0.01  # 20周为基准，每早1周增加1%权重

    # 新增：计算通过概率
    # 使用组内平均特征计算通过概率
    avg_height = group_data['身高'].mean()
    avg_weight = group_data['体重'].mean()
    avg_age = group_data['年龄'].mean()
    avg_bmi = group_data['孕妇BMI'].mean()
    # 在calculate_group_score函数中添加
    if len(valid_samples) > 1:
        预测误差标准差 = valid_samples['Y染色体浓度'].std()
    else:
        预测误差标准差 = 0.01  # 默认值
    pass_prob = calculate_pass_probability(
        avg_height, avg_weight, avg_age, avg_bmi, t_value, 
        预测误差标准差=预测误差标准差  # 假设预测误差标准差为0.01
    )
    
    # 新增约束条件：通过概率必须大于0.8
    if pass_prob <= 0.8:
        return 0.0  # 不满足约束条件，得分为0
    
    
    # # 计算总分：权重 * 达标比例 * 时间因子（强化早检测）
    # return time_weight * 达标比例 * time_factor
    # 对达标比例过低的分组增加提升空间
    if 达标比例 < 0.6:
        # 鼓励寻找能提高达标比例的稍晚孕周
        time_factor *= (0.6 / 达标比例) ** 0.3 if 达标比例!=0 else 0  # 适度调整
        
    # 新增风险分散度惩罚项
    avg_bmi = group_data['孕妇BMI'].mean()
    risk_level = get_risk_level(t_value, avg_bmi)
    
    # 风险等级权重 (鼓励选择更低风险的方案)
    risk_weights = {'低': 1.2, '中': 1.0, '较高': 0.8, '高': 0.6}
    risk_factor = risk_weights.get(risk_level, 1.0)
    
    # 最终得分 = 原始得分 * 风险因子 * 时间分散因子
    time_dispersion_factor = 1.0  # 可以通过历史推荐时间计算
    
    return time_weight * 达标比例 * time_factor * risk_factor * time_dispersion_factor

# ---------------------------
# 单分组优化函数
# ---------------------------
# 在优化单个分组时增加多样化搜索
def optimize_single_group(group_data: pd.DataFrame) -> float:
    # 多次运行优化并选择最佳结果，增加多样性
    best_score = -np.inf
    # 多次不同随机种子的搜索
    for seed in [42, 43, 44, 45]:
        res = gp_minimize(
            lambda t: -calculate_group_score(group_data, t),
            [Real(10, 23)],
            n_calls=15,  # 增加迭代次数
            random_state=seed,
            acq_func='EI',
            verbose=False,
            n_jobs=1
        )
        current_score = -res.fun
        if current_score > best_score:
            best_score = current_score
    return best_score

# ---------------------------
# 改进的优化算法（均衡样本量）
# ---------------------------
def optimize_bmi_groups(
    data: pd.DataFrame, 
    max_groups: int = 5,
    min_samples: int = 50,  # 提高最小样本量要求
    ideal_group_size: int = 150  # 理想的分组样本量
) -> List[float]:
    data = data.sort_values('孕妇BMI').reset_index(drop=True)
    n = len(data)
    
    # 计算建议的分组数，使样本量更均衡
    suggested_groups = max(2, min(max_groups, n // ideal_group_size))
    print(f"为均衡样本量，使用{suggested_groups}个分组（理想每组{ideal_group_size}样本）")
    
    if n < min_samples * suggested_groups:
        suggested_groups = max(2, n // min_samples)
        print(f"数据量有限，调整为{suggested_groups}个分组")
    
    # 动态步长策略
    step = max(2, n // (suggested_groups * 10))  # 确保足够的评估点
    sample_indices = list(range(0, n, step)) + [n-1]
    n_reduced = len(sample_indices)
    print(f"原始样本数: {n}，优化后评估点: {n_reduced}")
    
    # DP表初始化
    dp = [[-np.inf] * (n_reduced+1) for _ in range(suggested_groups+1)]
    dp[0][0] = 0
    splits = [[[] for _ in range(n_reduced+1)] for _ in range(suggested_groups+1)]
    
    # 预处理评分矩阵
    score_matrix = np.zeros((n_reduced, n_reduced))
    total_pairs = 0
    group_list = []
    
    # max_span = min(20, n_reduced // suggested_groups * 2)
    # 在预处理阶段增加更严格的过滤条件
    max_span = min(15, n_reduced // suggested_groups * 2)  # 从20减少到15
    min_group_size = min_samples  # 从0.8*min_samples提高到min_samples
    max_group_size = ideal_group_size  # 从1.5*ideal_group_size降到ideal_group_size
    for i in range(n_reduced):
        for j in range(i, min(i + max_span, n_reduced)):
            start_idx = sample_indices[i]
            end_idx = sample_indices[j]
            group_size = end_idx - start_idx + 1
            
            # 更严格的样本量过滤
            if min_group_size <= group_size <= max_group_size:
                group_data = data.iloc[start_idx:end_idx+1]
                group_list.append((i, j, group_data))
                total_pairs += 1
    # for i in range(n_reduced):
    #     # 限制分组大小，避免样本量差异过大
    #     for j in range(i, min(i + max_span, n_reduced)):
    #         start_idx = sample_indices[i]
    #         end_idx = sample_indices[j]
    #         group_size = end_idx - start_idx + 1
            
    #         # 过滤样本量过小或过大的分组
    #         if min_samples * 0.8 <= group_size <= ideal_group_size * 1.5:
    #             group_data = data.iloc[start_idx:end_idx+1]
    #             group_list.append((i, j, group_data))
    #             total_pairs += 1
    
    print(f"需要计算的分组数: {total_pairs}")
    
    # 并行计算
    with tqdm(total=total_pairs, desc="计算分组评分矩阵（并行）") as pbar:
        results = Parallel(
            n_jobs=multiprocessing.cpu_count()-1,
            verbose=0,
            backend='loky'
        )(delayed(optimize_single_group)(group_data) for i, j, group_data in group_list)
        
        for idx, (i, j, _) in enumerate(group_list):
            score_matrix[i][j] = results[idx]
            pbar.update(1)
    from functools import lru_cache
    @lru_cache(maxsize=1000)
    def cached_calculate_best_time(start_idx, end_idx):
        group_data = data.iloc[start_idx:end_idx+1]
        return calculate_best_time(group_data)

    def calculate_best_time(group_data: pd.DataFrame) -> float:
        """计算单个分组的最佳检测时间"""
        res = gp_minimize(
            lambda t: -calculate_group_score(group_data, t),
            [Real(10, 23)],
            n_calls=15,
            random_state=42,
            verbose=False
        )
        return round(res.x[0], 1)
    # 动态规划（加入样本量均衡惩罚）
    with tqdm(total=suggested_groups*n_reduced, desc="动态规划优化分组") as pbar:
        # 新增：记录每个状态的最佳时间
        best_times = [[None for _ in range(n_reduced+1)] for _ in range(suggested_groups+1)]
    

        for k in range(1, suggested_groups+1):
            for i in range(1, n_reduced+1):
                min_j = max(k-1, i - max_span)
                # 只检查最有希望的j值（按score_matrix排序）
                candidates = [(j, score_matrix[j][i-1]) for j in range(min_j, i)]
                candidates.sort(key=lambda x: -x[1])  # 按评分降序

                # for j in range(min_j, i):
                for j, _ in candidates[:5]:  # 只检查前5个候选
                    # group_data = data.iloc[sample_indices[j]:sample_indices[i-1]+1]
                    # current_time = calculate_best_time(group_data)
                    # 在动态规划中使用缓存版本
                    current_time = cached_calculate_best_time(sample_indices[j], sample_indices[i-1])
                    
                    if k > 1 and best_times[k-1][j] and abs(current_time - best_times[k-1][j]) < 1.0:
                        continue

                    # 在动态规划循环中添加
                    if k > 1 and dp[k-1][i] < 0.1:  # 如果前一组的得分已经很低
                        break  # 提前终止后续分组
                    
                    current_score = dp[k-1][j] + score_matrix[j][i-1]
                    
                    # 计算当前分组大小
                    start_idx = sample_indices[j] if j < len(sample_indices) else 0
                    end_idx = sample_indices[i-1] if (i-1) < len(sample_indices) else n-1
                    group_size = end_idx - start_idx + 1
                    
                    # 对偏离理想样本量的分组进行惩罚
                    size_ratio = group_size / ideal_group_size
                    if size_ratio < 0.5 or size_ratio > 2.0:  # 样本量偏差过大
                        current_score *= 0.7  # 显著降低分数
                    elif size_ratio < 0.7 or size_ratio > 1.5:  # 样本量偏差较大
                        current_score *= 0.9  # 适当降低分数
                    
                    if current_score > dp[k][i]:
                        dp[k][i] = current_score
                        splits[k][i] = splits[k-1][j] + [j]
                        best_times[k][i] = current_time  # 记录最佳时间
                        
                pbar.update(1)
    
    # 提取分割点
    best_splits = splits[suggested_groups][n_reduced]
    # 备选方案：若自动分组失败，使用均衡样本量的分位数分组
    if len(best_splits) < suggested_groups - 1:
        print("使用均衡样本量的分位数分组方案")
        # 按样本量而非值范围进行分位数分割
        quantiles = np.linspace(0, n, suggested_groups+1, dtype=int)[1:-1]
        bounds = [data.iloc[q]['孕妇BMI'] for q in quantiles]
        bounds = [data['孕妇BMI'].min()] + bounds + [data['孕妇BMI'].max()]
    else:
        bounds = [data['孕妇BMI'].min()]
        for split in best_splits:
            if 0 <= split < len(sample_indices):
                bounds.append(data.iloc[sample_indices[split]]['孕妇BMI'])
        bounds.append(data['孕妇BMI'].max())
    
    # 去重并确保至少2个分组
    bounds = sorted(list(set(bounds)))
    if len(bounds) < 3:  # 至少需要2个分组（3个边界）
        mid = (bounds[0] + bounds[-1]) / 2
        bounds = [bounds[0], mid, bounds[-1]]
    
    return bounds

# 改进风险等级划分逻辑
def get_risk_level(week: float, bmi: float, 
                  bmi_quantiles: dict = None, 
                  week_quantiles: dict = None) -> str:
    """
    基于数据分位数的动态风险划分
    参数:
        week: 孕周
        bmi: BMI值
        bmi_quantiles: BMI分位数字典 {'q20': 26.5, 'q40': 30.0, ...}
        week_quantiles: 孕周分位数字典
    返回:
        风险等级 ('低'/'中'/'高')
    """
    # 默认分位数（可根据实际数据计算）
    default_bmi_quantiles = {
        'q20': 26.5, 'q40': 30.0, 'q60': 32.5, 'q80': 34.2
    }
    default_week_quantiles = {
        'q20': 16.0, 'q40': 18.0, 'q60': 19.5, 'q80': 21.0
    }
    
    bmi_q = bmi_quantiles or default_bmi_quantiles
    week_q = week_quantiles or default_week_quantiles
    
    # 动态风险判断
    if week <= week_q['q20']:
        if bmi < bmi_q['q20']: return '低'
        elif bmi < bmi_q['q60']: return '中'
        else: return '高'
    elif week <= week_q['q40']:
        if bmi < bmi_q['q20']: return '低'
        elif bmi < bmi_q['q40']: return '中'
        elif bmi < bmi_q['q80']: return '较高'
        else: return '高'
    else:
        if bmi < bmi_q['q40']: return '中'
        else: return '高'

# 新增：临床优先级计算函数
def clinical_priority(bmi: float, week: float) -> float:
    """
    计算临床优先级得分（综合BMI和孕周因素）
    逻辑：
    1. BMI越高，检测难度越大，优先级越高
    2. 孕周过早/过晚时，优先级适当提高（需平衡检测可行性）
    """
    # 1. BMI基础优先级（非线性增长，高BMI增幅更大）
    if bmi < 25:
        bmi_priority = 1.0  # 正常BMI，基础优先级
    elif 25 <= bmi < 30:
        bmi_priority = 1.0 + (bmi - 25) * 0.08  # 轻度超重，缓慢提升
    elif 30 <= bmi < 35:
        bmi_priority = 1.4 + (bmi - 30) * 0.12  # 中度超重，加速提升
    else:
        bmi_priority = 2.0 + (bmi - 35) * 0.15  # 重度超重，大幅提升
    bmi_priority = np.clip(bmi_priority, 1.0, 3.0)  # 限制优先级范围
    
    # 2. 孕周修正因子（避免过早/过晚检测，同时兼顾可行性）
    if 14 <= week <= 18:  # 最佳检测窗口，无需修正
        week_factor = 1.0
    elif 12 <= week < 14 or 18 < week <= 20:  # 次优窗口，轻微修正
        week_factor = 1.1
    else:  # 边缘窗口（10-12周或20-23周），需要更高优先级推动检测
        week_factor = 1.2
    
    # 综合优先级（BMI为主，孕周为辅）
    return bmi_priority * week_factor

def adjust_recommendation_times(results: pd.DataFrame, male_data) -> pd.DataFrame:
    """调整推荐时间确保足够分散"""
    # 检查是否需要调整
    times = results['最佳检测孕周'].values
    if len(times) <= 1:
        return results  # 单个分组无需调整
    
    avg_interval = (max(times) - min(times)) / (len(times)-1)
    
    # 只有当时间点过于集中时才调整
    if avg_interval < 1.0:
        print(f"检测时间过于集中(平均间隔{avg_interval:.1f}周)，进行分散调整")
        
        # 计算新的等间隔时间点
        min_time = max(10, min(times)-1)  # 确保不低于10周
        max_time = min(23, max(times)+1)  # 确保不超过23周
        new_times = np.linspace(min_time, max_time, len(times))
        
        # 应用调整
        results['最佳检测孕周'] = np.round(new_times, 1)
        
        # 更新相关指标
        for i, row in results.iterrows():
            # 更新风险等级
            bmi_range = row['BMI范围']
            lower, upper = map(float, bmi_range.split('-'))
            avg_bmi = (lower + upper) / 2
            results.at[i, '风险等级'] = get_risk_level(row['最佳检测孕周'], avg_bmi)
            
            # 重新计算达标比例（可选）
            group_data = male_data[
                (male_data['孕妇BMI'] >= lower) & 
                (male_data['孕妇BMI'] < upper)
            ]
            valid_samples = group_data[group_data['Y染色体浓度'] > 0.04]
            early_records = valid_samples[valid_samples['检测孕周'] <= row['最佳检测孕周']]
            results.at[i, '达标比例'] = len(early_records)/len(valid_samples) if len(valid_samples)>0 else 0
    
    return results

# ---------------------------
# 主流程
# ---------------------------
def main(df: pd.DataFrame) -> pd.DataFrame:
    print("\n=== 数据预处理 ===")
    male_data = preprocess_data(df)
    print(f"有效男胎样本数: {len(male_data)}")

    # 计算实际数据分位数
    # male_data = preprocess_data(df)
    bmi_quantiles = {
        'q20': male_data['孕妇BMI'].quantile(0.2),
        'q40': male_data['孕妇BMI'].quantile(0.4),
        'q60': male_data['孕妇BMI'].quantile(0.6),
        'q80': male_data['孕妇BMI'].quantile(0.8)
    }
    week_quantiles = {
        'q20': male_data['检测孕周'].quantile(0.2),
        'q40': male_data['检测孕周'].quantile(0.4),
        'q60': male_data['检测孕周'].quantile(0.6),
        'q80': male_data['检测孕周'].quantile(0.8)
    }
    
    
    # 检查BMI分布
    bmi_stats = male_data['孕妇BMI'].describe()
    print("BMI分布统计:")
    print(bmi_stats)
    
    print("\n=== 模型构建 ===")
    time_model, poly, scaler = build_time_prediction_model(male_data)
    
    print("\n=== BMI分组优化 ===")
    # 计算理想分组数（基于总样本量）
    ideal_group_count = max(2, min(5, len(male_data) // 150))  # 每组约150样本
    bmi_bounds = optimize_bmi_groups(
        male_data,
        max_groups=ideal_group_count,
        ideal_group_size=150  # 目标每组150样本
    )
    print(f"获得BMI分组边界: {bmi_bounds}")
    
    # 并行计算各组最佳检测时间
    groups = []
    group_list = []
    
    # 准备所有分组数据
    for i in range(len(bmi_bounds)-1):
        lower, upper = bmi_bounds[i], bmi_bounds[i+1]
        group_data = male_data[
            (male_data['孕妇BMI'] >= lower) & 
            (male_data['孕妇BMI'] < upper)
        ]
        # 调整分组边界，确保最小样本量
        if len(group_data) < 50 and i > 0:
            print(f"调整分组边界，确保样本量充足 (原{len(group_data)}样本)")
            # 合并上一分组
            lower = bmi_bounds[i-1]
            group_data = male_data[
                (male_data['孕妇BMI'] >= lower) & 
                (male_data['孕妇BMI'] < upper)
            ]
            # 从列表中移除上一分组
            if groups:
                groups.pop()
        
        if len(group_data) >= 50:  # 确保足够样本量
            group_list.append((i, lower, upper, group_data))
    
    # 并行优化每个分组的最佳时间
    with tqdm(total=len(group_list), desc="计算各组最佳检测时间（并行）") as pbar:
        results = Parallel(
            n_jobs=multiprocessing.cpu_count()-1,
            verbose=0,
            backend='loky'
        )(delayed(optimize_single_group)(group_data) for i, lower, upper, group_data in group_list)
        
        # 整理结果
        for idx, (i, lower, upper, group_data) in enumerate(group_list):
            best_score = round(results[idx], 3)
            
            # 重新计算最佳时间，特别关注更早的孕周
            res = gp_minimize(
                lambda t: -calculate_group_score(group_data, t),
                [Real(10, 23)],  # 限制在23周以内
                n_calls=15,  # 增加迭代次数以找到更精确的早孕周
                random_state=42,
                verbose=False
            )
            best_time = round(res.x[0], 1)
            time_weight = get_time_weight(best_time)
            
            valid_samples = group_data[group_data['Y染色体浓度'] > 0.04]
            early_records = valid_samples[valid_samples['检测孕周'] <= best_time]
            达标比例 = round(len(early_records)/len(valid_samples), 3) if len(valid_samples) > 0 else 0
            
            groups.append({
                'BMI范围': f"{lower:.1f}-{upper:.1f}",
                '最佳检测孕周': best_time,
                '时间权重': time_weight,
                '达标比例': 达标比例,
                '总分': best_score,
                '样本量': len(group_data),
                # '风险等级': get_risk_level(best_time, (lower+upper)/2)  # 使用分组平均BMI
                '风险等级': get_risk_level(best_time, (lower+upper)/2, bmi_quantiles, week_quantiles)
            })
            pbar.update(1)

    # 结果展示
    results = pd.DataFrame(groups)

    print("\n=== 结果后处理 ===")
    # optimal_groups = adjust_recommendation_times(results)
    # 3. 调整时间时显式传递数据
    optimal_groups = adjust_recommendation_times(
        results, 
        male_data  # 明确传递所需数据
    )
    
    # 结果展示
    print("\n=== 最佳分组策略 ===")
    print(optimal_groups.to_markdown(index=False))

    plot_results(optimal_groups)
    return optimal_groups

def plot_results(results: pd.DataFrame) -> None:
    plt.figure(figsize=(14, 7))
    
    color_map = {'低': '#4CAF50', '中': '#FFC107', '高': '#F44336'}
    bar_width = 0.6
    
    for _, group in results.iterrows():
        plt.bar(
            group['BMI范围'], 
            group['最佳检测孕周'],
            width=bar_width,
            label=f"{group['风险等级']}风险",
            # color=color_map[group['风险等级']],
            alpha=0.8,
            edgecolor='black'
        )
        
        plt.text(
            group['BMI范围'], 
            group['最佳检测孕周'] + 0.3,
            f"{group['最佳检测孕周']}周\n总分:{group['总分']:.2f}\n样本:{group['样本量']}",
            ha='center',
            va='bottom',
            fontsize=10,
            bbox=dict(facecolor='white', alpha=0.8, edgecolor='none', pad=0.5)
        )
    
    plt.xlabel('BMI分组', fontsize=12, labelpad=10)
    plt.ylabel('推荐检测孕周', fontsize=12, labelpad=10)
    plt.title('基于BMI的最佳NIPT检测时间推荐', fontsize=14, pad=20, fontweight='bold')
    
    plt.legend(title='风险等级', loc='upper right', framealpha=0.9)
    plt.grid(axis='y', linestyle='--', alpha=0.5)
    
    plt.ylim(9, 24)  # 调整y轴范围，聚焦10-23周
    
    plt.tight_layout()
    plt.savefig( os.path.join(out_dir, 'ques3_optimal_bmi_groups.png'), dpi=300, bbox_inches='tight')
    plt.close()
    print("\n已保存可视化结果: ques3_optimal_bmi_groups.png")

# ---------------------------
# 敏感性分析
# ---------------------------
def sensitivity_analysis(results: pd.DataFrame, male_data: pd.DataFrame, 
                       error_rates: List[float] = [0.01, 0.05, 0.1]) -> pd.DataFrame:
    """分析检测误差对结果的影响
    
    Args:
        results: 主流程的结果DataFrame
        male_data: 预处理后的男胎数据
        error_rates: 要分析的误差率列表
        
    Returns:
        敏感性分析结果DataFrame
    """
    sensitivity_results = []
    
    for error_rate in error_rates:
        for _, group in results.iterrows():
            lower, upper = map(float, group['BMI范围'].split('-'))
            group_data = male_data[
                (male_data['孕妇BMI'] >= lower) & 
                (male_data['孕妇BMI'] < upper)
            ]
            
            if len(group_data) == 0:
                continue
                
            # 模拟误差影响：达标比例按误差率下降
            valid_samples = group_data[group_data['Y染色体浓度'] > 0.04]
            early_records = valid_samples[valid_samples['检测孕周'] <= group['最佳检测孕周']]
            original_rate = len(early_records) / len(valid_samples) if len(valid_samples) > 0 else 0
            adjusted_rate = original_rate * (1 - error_rate)
            
            # 计算调整后的总分
            adjusted_score = group['时间权重'] * adjusted_rate
            
            sensitivity_results.append({
                '误差率': f"{error_rate*100:.0f}%",
                'BMI范围': group['BMI范围'],
                '原始总分': round(group['总分'], 3),
                '调整后总分': round(adjusted_score, 3),
                '总分损失': round(group['总分'] - adjusted_score, 3)
            })
    
    sensitivity_df = pd.DataFrame(sensitivity_results)
    
    # 可视化
    plt.figure(figsize=(14, 7))
    for error_rate in error_rates:
        subset = sensitivity_df[sensitivity_df['误差率'] == f"{error_rate*100:.0f}%"]
        plt.plot(subset['BMI范围'], subset['调整后总分'], 
                marker='o', markersize=8, linewidth=2,
                label=f'误差率={error_rate*100:.0f}%')
    
    plt.xlabel('BMI分组', fontsize=12)
    plt.ylabel('调整后总分', fontsize=12)
    plt.title('检测误差对总分的影响', fontsize=14, pad=20)
    plt.legend(title='误差率', loc='upper right')
    plt.grid(True, linestyle='--', alpha=0.7)
    plt.ylim(0, 2.1)  # 最大可能得分为2.0
    plt.tight_layout()
    plt.savefig( os.path.join(out_dir, 'ques3_sensitivity_analysis.png'), dpi=300)
    plt.close()
    print("\n已保存敏感性分析结果: sensitivity_analysis.png")
    
    return sensitivity_df

# ==================== 改进的检测误差分析部分 ====================
def improved_error_analysis(optimal_groups, male_data):
    """优化后完整误差分析方法，增强结果梯度与临床合理性"""
    # 1. 临床参数初始化 - 精细化误差范围
    sigma_range = np.linspace(0.04, 0.25, 12)   # 增加采样点，提升梯度细腻度
    sigma_multipliers = [1.0, 1.5, 2.0, 2.5, 2.8, 3.0]  # 增加2.8倍数，细化警告触发点
    
    # 计算各BMI分组的Y染色体浓度标准差（增强连续性差异）
    base_std_list = []
    for _, group in optimal_groups.iterrows():
        lower, upper = map(float, group['BMI范围'].split('-'))
        group_data = male_data[
            (male_data['孕妇BMI'] >= lower) & 
            (male_data['孕妇BMI'] < upper)
        ]
        valid_samples = group_data[group_data['Y染色体浓度'] > 0.04]
        if len(valid_samples) > 1:
            # 按BMI区间连续增强标准差（每增加1单位BMI，标准差提升2%）
            bmi_factor = 1.0 + (max(lower, 25) - 25) * 0.02
            base_std = valid_samples['Y染色体浓度'].std() * bmi_factor
            base_std_list.append(max(base_std, 0.04))
        else:
            # 按BMI区间设置连续递增的默认标准差
            bmi_mid = (lower + upper) / 2
            base_std_list.append(0.04 + (bmi_mid - 25) * 0.005)
    
    def calc_pass_probability(height, weight, age, bmi, week, sigma):
        """精细化概率计算，增强分组梯度差异"""
        # 1. BMI对误差的连续增强模型（每增加1单位BMI，误差因子提升3%-5%）
        if bmi < 25:
            bmi_error_factor = 1.0 + (bmi - 20) * 0.01
        elif bmi < 30:
            bmi_error_factor = 1.1 + (bmi - 25) * 0.09
        elif bmi < 35:
            bmi_error_factor = 1.5 + (bmi - 30) * 0.15
        else:
            bmi_error_factor = 2.2 + (bmi - 35) * 0.22  # 最高BMI组增强更显著
        
        # 2. 孕周对误差的非线性影响
        week_error_factor = 1.0
        if week < 13:
            week_error_factor = 1.2 + (12 - week) * 0.12
        elif week > 21:
            week_error_factor = 1.05 + (week - 21) * 0.08
        
        adjusted_sigma = sigma * bmi_error_factor * week_error_factor
        
        # 3. 计算基础达标概率
        base_prob = calculate_pass_probability(
            height, weight, age, bmi, week, adjusted_sigma
        )
        
        # 4. 分组特异性误差衰减（核心优化：增强连续性差异）
        if bmi < 30:
            # 低BMI组衰减较慢，保留较高概率
            error_decay = 1.0 - (adjusted_sigma ** 0.85) * 0.75
        elif bmi < 35:
            # 中BMI组中等衰减
            error_decay = 1.0 - (adjusted_sigma ** 0.85) * 1.1
        else:
            # 高BMI组衰减更显著，拉开与中BMI组差距
            error_decay = 1.0 - (adjusted_sigma ** 0.85) * 1.5
        
        # 确保衰减不过度，同时保留差异
        error_decay = max(0.62 if bmi >=35 else 0.65 if bmi >=30 else 0.78, error_decay)
        
        # 5. 分组特异性孕周修正
        week_factor = 1.0
        if bmi < 30:
            if week < 12:
                week_factor = 0.8 + (week - 10) * 0.08
            elif week > 20:
                week_factor = 0.92 - (week - 20) * 0.02
        else:
            if week < 12:
                week_factor = 0.65 + (week - 10) * 0.06
            elif week > 20:
                week_factor = 0.8 - (week - 20) * 0.04
        
        # 综合概率计算
        combined_prob = base_prob * week_factor * error_decay
        
        # 分组特异性下限（拉开最高组与中间组差距）
        if bmi < 30:
            return np.clip(combined_prob, 0.78, None)  # 低BMI组：78%下限
        elif bmi < 35:
            return np.clip(combined_prob, 0.65, None)  # 中BMI组：65%下限
        else:
            return np.clip(combined_prob, 0.62, None)  # 高BMI组：62%下限

    sensitivity_results = []
    robustness_results = []
    bmi_ranges = [(float(bmi_range.split('-')[0]), float(bmi_range.split('-')[1])) 
                 for bmi_range in optimal_groups['BMI范围']]
    n_segments = len(bmi_ranges)
    
    # 3. 敏感性分析
    plt.figure(figsize=(12, 6))
    for i, bmi_range in enumerate(bmi_ranges):
        lower, upper = bmi_range
        avg_bmi = (lower + upper) / 2
        group_data = male_data[
            (male_data['孕妇BMI'] >= lower) & 
            (male_data['孕妇BMI'] < upper)
        ]
        avg_height = group_data['身高'].mean()
        avg_weight = group_data['体重'].mean()
        avg_age = group_data['年龄'].mean()
        optimal_week = optimal_groups.iloc[i]['最佳检测孕周']
        
        pass_probs = []
        for sigma in sigma_range:
            pass_prob = calc_pass_probability(
                avg_height, avg_weight, avg_age, avg_bmi, optimal_week, sigma
            )
            pass_probs.append(pass_prob)
            
            sensitivity_results.append({
                'bmi_segment': i,
                'bmi_range': f"{lower:.1f}-{upper:.1f}",
                'sigma': sigma,
                'pass_probability': pass_prob,
                'optimal_week': optimal_week
            })
        
        plt.plot(sigma_range, pass_probs, 
                 label=f'BMI分段 {lower:.1f}-{upper:.1f}',
                 marker='o')

    plt.axhline(y=0.8, color='r', linestyle='--', label='达标阈值 (0.8)')
    plt.xlabel('Y染色体浓度误差标准差', fontsize=12)
    plt.ylabel('达标概率', fontsize=12)
    plt.title('不同误差水平下的达标概率变化', fontsize=14)
    plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left')
    plt.grid(True, alpha=0.3)
    plt.tight_layout()
    plt.savefig(os.path.join(out_dir, 'improved_sensitivity_analysis.png'), dpi=300, bbox_inches='tight')
    plt.close()

    # 4. 稳健性分析（优化时间调整范围）
    for idx, sigma_mult in enumerate(sigma_multipliers):
        # 按BMI连续增强实际误差（高BMI组误差增长更快）
        modified_std = [base_std * (sigma_mult * (1 + i*0.1)) 
                       for i, base_std in enumerate(base_std_list)]
        
        for i in range(n_segments):
            lower, upper = bmi_ranges[i]
            avg_bmi = (lower + upper) / 2
            group_data = male_data[
                (male_data['孕妇BMI'] >= lower) & 
                (male_data['孕妇BMI'] < upper)
            ]
            avg_height = group_data['身高'].mean()
            avg_weight = group_data['体重'].mean()
            avg_age = group_data['年龄'].mean()
            original_week = optimal_groups.iloc[i]['最佳检测孕周']
            
            # 连续化最大时间调整限制（随BMI递增）
            max_adjustment = 1.5 + (avg_bmi - 25) * 0.15
            max_adjustment = min(max_adjustment, 3.0)  # 上限3周
        
            # 连续化时间惩罚系数（随BMI递减）
            time_penalty_coef = 0.5 - (avg_bmi - 25) * 0.015
            time_penalty_coef = max(time_penalty_coef, 0.2)  # 下限0.2
        
            def segment_objective(x):
                week = x[0]
                # 硬性约束
                if not (10 <= week <= 23):
                    return 1e8
                
                # 限制最大时间调整幅度
                time_deviation = abs(week - original_week)
                if time_deviation > max_adjustment:
                    return 1e7 + (time_deviation - max_adjustment) * 1e6
                
                # 连续化达标阈值（随BMI递减）
                threshold = 0.88 - (avg_bmi - 25) * 0.008
                threshold = max(threshold, 0.80)  # 下限0.80
                
                prob = calc_pass_probability(
                    avg_height, avg_weight, avg_age, avg_bmi, week, modified_std[i]
                )
                if prob < threshold:
                    return 1e5 + (threshold - prob) * 1e4
                
                # 临床优先级和时间权重
                priority = clinical_priority(avg_bmi, week)
                time_weight = get_time_weight(week)
                base_score = time_weight * priority
                
                # 时间调整惩罚（分段惩罚）
                if time_deviation < max_adjustment * 0.5:
                    deviation_penalty = (time_deviation **1.3) * time_penalty_coef * 0.3
                elif time_deviation < max_adjustment * 0.8:
                    deviation_penalty = (time_deviation** 1.8) * time_penalty_coef * 0.6
                else:
                    deviation_penalty = (time_deviation ** 2.5) * time_penalty_coef * 1.2
                
                final_score = base_score - deviation_penalty
                return -final_score
            
            # PSO参数
            x_opt, _ = pso(
                segment_objective,
                lb=[10], ub=[23],
                swarmsize=80,
                maxiter=150,
                omega=0.7,
                phip=1.5,
                phig=1.5,
                minfunc=1e-4,
                debug=False
            )
            
            # 记录结果
            adjustment_category = "小幅" if abs(x_opt[0] - original_week) <= 1 else \
                                "中度" if abs(x_opt[0] - original_week) <= 2 else "大幅"
            
            robustness_results.append({
                'sigma_multiplier': sigma_mult,
                'actual_sigma': modified_std[i],
                'bmi_segment': i,
                'bmi_range': f"{lower:.1f}-{upper:.1f}",
                'original_week': original_week,
                'adjusted_week': x_opt[0],
                'time_change': x_opt[0] - original_week,
                'adjustment_category': adjustment_category,
                'pass_probability': calc_pass_probability(
                    avg_height, avg_weight, avg_age, avg_bmi, x_opt[0], modified_std[i]
                ),
                'original_prob': calc_pass_probability(
                    avg_height, avg_weight, avg_age, avg_bmi, original_week, modified_std[i]
                )
            })

    # 5. 保存结果
    sensitivity_df = pd.DataFrame(sensitivity_results)
    robustness_df = pd.DataFrame(robustness_results)
    
    sensitivity_df.to_csv(os.path.join(out_dir, 'improved_sensitivity_results.csv'), index=False)
    robustness_df.to_csv(os.path.join(out_dir, 'improved_robustness_results.csv'), index=False)

    # 6. 可视化稳健性分析结果
    plt.figure(figsize=(14, 8))
    
    # 时间变化图
    plt.subplot(2, 1, 1)
    for i in range(n_segments):
        seg_data = robustness_df[robustness_df['bmi_segment'] == i]
        plt.plot(seg_data['sigma_multiplier'], seg_data['time_change'],
                 marker='o', label=f'BMI分段 {bmi_ranges[i][0]:.1f}-{bmi_ranges[i][1]:.1f}')
    plt.axhline(y=0, color='k', linestyle='--')
    plt.xlabel('误差倍数')
    plt.ylabel('时间变化(周)')
    plt.title('误差对检测时间的影响')
    plt.legend()
    plt.grid(True, alpha=0.3)

    # 达标概率图
    plt.subplot(2, 1, 2)
    for i in range(n_segments):
        seg_data = robustness_df[robustness_df['bmi_segment'] == i]
        plt.plot(seg_data['sigma_multiplier'], seg_data['pass_probability'],
                 marker='o', label=f'BMI分段 {bmi_ranges[i][0]:.1f}-{bmi_ranges[i][1]:.1f}')
    plt.axhline(y=0.8, color='r', linestyle='--', label='达标阈值')
    plt.xlabel('误差倍数')
    plt.ylabel('达标概率')
    plt.title('误差对达标概率的影响')
    plt.legend()
    plt.grid(True, alpha=0.3)

    plt.tight_layout()
    plt.savefig(os.path.join(out_dir, 'improved_robustness_analysis.png'), dpi=300, bbox_inches='tight')
    plt.close()

    # 7. 输出关键发现
    print("\n=== 改进的误差分析关键发现 ===")
    for i in range(n_segments):
        seg_data = robustness_df[robustness_df['bmi_segment'] == i]
        max_change = seg_data['time_change'].abs().max()
        min_prob = seg_data['pass_probability'].min()
        
        # 敏感度评估
        if max_change > 2.0:
            sensitivity = '高'
        elif max_change > 1.0:
            sensitivity = '中'
        else:
            sensitivity = '低'
        
        print(f"BMI分段 {bmi_ranges[i][0]:.1f}-{bmi_ranges[i][1]:.1f}:")
        print(f"  最大时间变化: {max_change:.2f} 周")
        print(f"  最低达标概率: {min_prob:.2%}")
        print(f"  对误差敏感度: {sensitivity}")
        
        # 检查是否存在不达标情况
        if min_prob < 0.8:
            # 找到触发警告的最小误差倍数
            min_sigma = seg_data.loc[seg_data['pass_probability'].idxmin(), 'sigma_multiplier']
            print(f"  警告: 在σ={min_sigma:.1f}时达标概率低于80%")



# 在主程序中添加调用
if __name__ == "__main__":
    try:
        print("=== 数据加载 ===")
        df = pd.read_csv('ques1_prepare_data/sheet1_男胎检测数据_fillnan_sifted_outlier.csv')
        
        if len(df) < 200:
            raise ValueError("样本量过小（<200），无法进行有效分组")
        
        print("\n=== 主优化流程 ===")
        optimal_groups = main(df)
        
        if len(optimal_groups) < 2:
            print("警告：分组数量不足，尝试增加样本量或调整分组参数")
        
        print("\n=== 敏感性分析 ===")
        male_data = preprocess_data(df)
        sensitivity_results = sensitivity_analysis(optimal_groups, male_data)
        print("\n敏感性分析结果：")
        print(sensitivity_results.to_markdown(index=False))
        
        # 新增：运行改进的误差分析
        print("\n=== 运行改进的误差分析 ===")
        improved_error_analysis(optimal_groups, male_data)
        
    except Exception as e:
        print(f"程序运行出错: {str(e)}")
        raise